# strategy.py
# -*- coding: utf-8 -*-
import time
from datetime import datetime, timedelta
import numpy as np
import math
import pandas as pd
import akshare as ak

# 假设 get_data.py 中封装了对底层数据源（akshare 或 joinquant）的适配函数：
#   get_price(symbol, start_date=None, end_date=None, count=None, frequency='daily', fields=['close',...'])
#   get_index_price(index_symbol, start_date=None, end_date=None)
#   get_security_info(symbol) -> object with display_name attribute
#   get_trade_days(end_date=None, count=..., start_date=None)
from get_data import get_price, get_index_price, get_trade_days
# import get_data

# --- 策略参数 ---
stock_pool = [
    '159915.XSHE', # 易方达创业板ETF
    '510300.XSHG', # 华泰柏瑞沪深300ETF
    '510500.XSHG', # 南方中证500ETF
]

stock_num = 1           # 买入评分最高的前 stock_num 只股票
momentum_day = 23       # 最新动量参考最近 momentum_day 的（调整为实际可获取天数）
ref_stock = '000300.XSHG'   # 用 ref_stock 做择时计算的基础数据（在 get_data 中适配）
N = 18 # RSRS 斜率回归窗口
M = 120 # 斜率序列用于计算 zscore 的窗口
score_threshold = 0.7 # rsrs 标准分阈值

mean_day = 20
mean_diff_day = 3

day = 1  # 用于 get_rank 的 end_date 偏移（保持原脚本设计）

# import akshare as ak

def get_security_info(code: str):
    """
    获取股票基本信息
    """
    # 处理股票代码格式
    if code.endswith('.XSHG'):
        stock_code = code.replace('.XSHG', '')
    elif code.endswith('.XSHE'):
        stock_code = code.replace('.XSHE', '')
    else:
        stock_code = code
        
    try:
        df = ak.stock_individual_info_em(symbol=stock_code)
        return df
    except Exception as e:
        print(f"获取股票信息失败 {code}: {e}")
        # 返回一个简单的对象
        class SimpleInfo:
            def __init__(self, code):
                self.display_name = code
        return SimpleInfo(code)


# -------------------------
# 统计/回归工具函数
# -------------------------
def get_ols(x, y):
    """
    对 x,y 一维 numpy 或 pandas.Series 做一元线性回归，返回 (intercept, slope, r2)
    使用标准 R^2 公式，若常数方差为0返回 r2=0 保护。
    """
    x = np.asarray(x)
    y = np.asarray(y)
    if len(x) < 2:
        return (0.0, 0.0, 0.0)
    # 直线拟合
    slope, intercept = np.polyfit(x, y, 1)
    y_pred = slope * x + intercept
    ss_res = np.sum((y - y_pred) ** 2)
    ss_tot = np.sum((y - np.mean(y)) ** 2)
    r2 = 1 - ss_res / ss_tot if ss_tot != 0 else 0.0
    return (intercept, slope, r2)

# -------------------------
# 选股（动量）模块
# -------------------------
def get_rank(stock_pool, end_date=None):
    """
    返回按动量评分排序后 top stock_num 的股票代码列表
    使用 get_price(...) 接口（假定它接受 end_date=datetime 或日期字符串，或 count 参数）
    """
    score_list = []
    # 计算基准日期
    if end_date is None:
        current_dt = datetime.now()
        end_date = current_dt - timedelta(days=day)
    # 避免超权限日期（更新到当前日期）
    if isinstance(end_date, datetime) and end_date > datetime.now():
        end_date = datetime.now()
    # 确保有足够的历史数据（至少需要30天）
    if isinstance(end_date, datetime):
        # 如果日期太新，回退到有足够数据的日子
        test_date = end_date - timedelta(days=30)
        if test_date.year < 2025 or (test_date.year == 2025 and test_date.month < 9):
            end_date = datetime(2025, 9, 20)  # 使用一个已知有数据的日期
    print(f"  动量计算基准日期: {end_date.strftime('%Y-%m-%d')}")
    for stock in stock_pool:
        try:
            # 这里假定 get_price 支持 end_date=datetime, count=int, frequency='daily'
            data = get_price(stock, end_date=end_date, count=momentum_day, frequency='daily', fields=['close'])
            if data is None or len(data) < 15:  # 最少需要15天数据进行有效计算
                print(f"  ⚠️  {stock} 数据不足: 需要至少15天，实际{0 if data is None else len(data)}天")
                score_list.append(0.0)
                continue
            elif len(data) < momentum_day:
                print(f"  📊 {stock} 使用可用数据: 目标{momentum_day}天，实际{len(data)}天")
            
            # 确保数据是纯numpy数组，完全避免pandas索引问题
            if isinstance(data, pd.DataFrame):
                if 'close' in data.columns:
                    close_values = data['close'].values
                else:
                    # 如果没有close列，尝试使用第一列
                    close_values = data.iloc[:, 0].values
            else:
                # 如果已经是numpy数组或其他类型，尝试直接转换
                close_values = np.asarray(data)
            
                
            # 使用对数价格做线性拟合得到 slope
            y = np.log(close_values)
            x = np.arange(len(y))
            slope, intercept = np.polyfit(x, y, 1)
            annualized_returns = math.pow(math.exp(slope), 250) - 1
            
            # 计算 R^2
            y_pred = slope * x + intercept
            ss_res = np.sum((y - y_pred) ** 2)
            ss_tot = np.sum((y - np.mean(y)) ** 2)
            r_squared = 1 - (ss_res / ss_tot) if ss_tot != 0 else 0.0
            score = annualized_returns * r_squared
            score_list.append(score)
            
            # 名称（假定 get_security_info 返回有 display_name 属性）
            info = get_security_info(stock)
            name = getattr(info, 'display_name', stock)
            print(f"  {name}({stock}): 年化收益={annualized_returns:.4f}, R²={r_squared:.4f}, 评分={score:.4f}")
        except Exception as e:
            print(f"  ❌ {stock} 计算失败: {e}")
            score_list.append(0.0)
    
    # 修复这里的索引问题 - 使用传入的stock_pool而不是g.stock_pool
    stock_dict = dict(zip(stock_pool, score_list))
    sort_list = sorted(stock_dict.items(), key=lambda item: item[1], reverse=True)
    rank_stock = [t[0] for t in sort_list[:min(stock_num, len(sort_list))]]
    return rank_stock

# -------------------------
# 择时（RSRS + MA）模块
# -------------------------
def initial_slope_series(end_date=None):
    """
    生成历史斜率序列（用于 RSRS zscore 计算）
    end_date: datetime 或 None（默认为 now）
    返回斜率列表（len >= 0）
    """
    if end_date is None:
        current_dt = datetime.now()
    else:
        current_dt = end_date
    # 取 M 天之前作为起点（用 calendar days 安全一点）
    start_dt = current_dt - timedelta(days=M + N + 30)  # 多取一些缓冲
    start_str = start_dt.strftime('%Y%m%d')
    end_str = current_dt.strftime('%Y%m%d')
    # 注意：这里我们假定 get_index_price(index_symbol, start_date, end_date) 返回 DataFrame 包含 'high'和'low'
    # 另外注意 index symbol 的格式应由 get_index_price 在 get_data 中统一处理（例如 '000300.XSHG' -> 'sh000300'）
    index_symbol = ref_stock  # 使用全局 ref_stock（例如 '000300.XSHG'）
    data = get_index_price(index_symbol, start_date=start_str, end_date=end_str)
    if data is None or len(data) < N:
        raise Exception("数据不足，无法生成初始斜率序列")
    # 选取最近 N+M 或全部可用
    if len(data) < N + M:
        required_data = data
    else:
        required_data = data.tail(N + M)
    slopes = []
    # 对每个可移动窗口计算 slope（使用 low 作为自变量，high 作为因变量，与原脚本一致）
    for i in range(0, len(required_data) - N + 1):
        low_series = required_data['low'].iloc[i:i+N].values
        high_series = required_data['high'].iloc[i:i+N].values
        _, slope, _ = get_ols(low_series, high_series)
        slopes.append(slope)
    return slopes

def get_zscore(slope_series):
    slope_series = np.asarray(slope_series)
    if len(slope_series) < 2:
        return 0.0
    mean = slope_series.mean()
    std = slope_series.std(ddof=0)
    if std == 0:
        return 0.0
    return (slope_series[-1] - mean) / std

def get_timing_signal(index_symbol, calculation_date=None):
    """
    返回 'BUY' / 'SELL' / 'KEEP'
    calculation_date: datetime 或 None（默认为 now）
    """
    global slope_series
    if calculation_date is None:
        current_dt = datetime.now()
    else:
        current_dt = calculation_date
    
    # 确保有足够的历史数据
    if current_dt > datetime(2025, 9, 23):
        current_dt = datetime(2025, 9, 23)  # 使用最后一个有数据的日期
    # MA 需要 recent mean_day + mean_diff_day 天的收盘价
    start_for_ma = (current_dt - timedelta(days=mean_day + mean_diff_day)).strftime('%Y%m%d')
    end_for_ma = current_dt.strftime('%Y%m%d')
    try:
        # 获取收盘价（使用 get_index_price 获取指数数据）
        close_data = get_index_price(index_symbol, start_date=start_for_ma, end_date=end_for_ma)
        if close_data is not None and 'close' in close_data.columns:
            close_data = close_data[['close']]
        if close_data is None or len(close_data) < mean_day + mean_diff_day:
            print(f"⚠️  MA数据不足: 需要{mean_day + mean_diff_day}天，实际{0 if close_data is None else len(close_data)}天")
            return "KEEP"
        # 计算当前 MA 和之前 MA（窗口分段）
        
        today_MA = close_data['close'].iloc[-mean_day:].mean()
        before_MA = close_data['close'].iloc[-(mean_day + mean_diff_day):-mean_diff_day].mean()
        # 获取 high/low 用于 RSRS（最近 N 天）
        start_for_rsrs = (current_dt - timedelta(days=N+5)).strftime('%Y%m%d')
        end_for_rsrs = current_dt.strftime('%Y%m%d')
        high_low_data = get_index_price(index_symbol, start_date=start_for_rsrs, end_date=end_for_rsrs)
        if high_low_data is not None and 'high' in high_low_data.columns and 'low' in high_low_data.columns:
            high_low_data = high_low_data[['high', 'low']]
        if high_low_data is None or len(high_low_data) < N:
            print(f"⚠️  RSRS数据不足: 需要{N}天，实际{0 if high_low_data is None else len(high_low_data)}天")
            return "KEEP"
        recent = high_low_data.tail(N)
        intercept, slope, r2 = get_ols(recent['low'].values, recent['high'].values)
        # 更新全局斜率序列
        slope_series.append(slope)
        # 动态计算 effective_M
        available_data = len(slope_series)
        effective_M = min(M, available_data) if available_data >= 30 else available_data
        if effective_M < 30:
            print(f"⚠️  斜率历史数据不足: 需要至少30个，实际{available_data}个")
            return "KEEP"
        recent_slopes = slope_series[-effective_M:]
        rsrs_score = get_zscore(recent_slopes) * r2
        print(f"  MA信号: 当前={today_MA:.4f}, 之前={before_MA:.4f}, 趋势={'上升' if today_MA > before_MA else '下降'}")
        print(f"  RSRS信号: 分数={rsrs_score:.4f}, 阈值=±{score_threshold}, R²={r2:.4f}, 使用{effective_M}天数据")
        ma_signal = today_MA > before_MA
        rsrs_buy_signal = rsrs_score > score_threshold
        rsrs_sell_signal = rsrs_score < -score_threshold
        if rsrs_buy_signal and ma_signal:
            return "BUY"
        elif rsrs_sell_signal and not ma_signal:
            return "SELL"
        else:
            return "KEEP"
    except Exception as e:
        print(f"❌ 择时信号计算失败: {e}")
        return "KEEP"

# 初始化斜率序列（在主运行前调用）
try:
    slope_series = initial_slope_series()
    if len(slope_series) > 0:
        slope_series = slope_series[:-1]  # 与原脚本逻辑保持一致，去掉最后一个以避免重复
    print(f"✅ 初始化斜率序列，长度={len(slope_series)}")
except Exception as e:
    print(f"⚠️  初始化斜率序列失败: {e}")
    slope_series = []

# -------------------------
# 主执行逻辑
# -------------------------
def run_single_day():
    global slope_series
    print("=== 今日策略运行 ===")
    try:
        # 重新初始化斜率序列（确保最新）
        try:
            slope_series = initial_slope_series()[:-1]
            print(f"✅ 斜率序列初始化完成，长度: {len(slope_series)}")
        except Exception as e:
            print(f"⚠️  重新初始化斜率序列失败: {e}")
        # 选股
        check_out_list = get_rank(stock_pool)
        print("\n📊 今日股票排名:")
        for i, stock_code in enumerate(check_out_list):
            info = get_security_info(stock_code)
            name = getattr(info, 'display_name', stock_code)
            print(f'  {i+1}. {name}({stock_code})')
        # 择时
        timing_signal = get_timing_signal(ref_stock)
        print(f'\n🎯 择时信号: {timing_signal}')
        if timing_signal == "BUY" and check_out_list:
            top_stock = check_out_list[0]
            info = get_security_info(top_stock)
            name = getattr(info, 'display_name', top_stock)
            print(f'💡 操作建议: 买入 {name}({top_stock})')
        elif timing_signal == "SELL":
            print('💡 操作建议: 卖出当前持仓')
        else:
            print('💡 操作建议: 保持当前仓位')
    except Exception as e:
        print(f'❌ 策略运行失败: {e}')
    print('=' * 50)

def quick_test():
    """
    快速测试最近几天的策略表现（使用 get_trade_days）
    """
    global slope_series
    print("=== 快速回测（最近5个交易日）===")
    try:
        end_date = datetime(2025, 6, 17)
        trade_days = get_trade_days(end_date=end_date, count=5)  # 假定返回可迭代的日期列表
        if not trade_days:
            print("没有交易日数据")
            return
        # 用第一个交易日初始化斜率序列
        init_date = trade_days[0]
        if isinstance(init_date, datetime):
            init_datetime = init_date
        else:
            init_datetime = datetime.combine(init_date, datetime.min.time())
        slope_series = initial_slope_series(init_datetime)[:-1]
        for trade_date in trade_days:
            if isinstance(trade_date, datetime):
                trade_dt = trade_date
            else:
                trade_dt = datetime.combine(trade_date, datetime.min.time())
            print(f"\n--- {trade_dt.strftime('%Y-%m-%d')} ---")
            check_out_list = get_rank(stock_pool, end_date=trade_dt)
            top_stock = check_out_list[0] if check_out_list else None
            if top_stock:
                info = get_security_info(top_stock)
                print(f"推荐股票: {getattr(info,'display_name',top_stock)}")
            timing_signal = get_timing_signal(ref_stock, calculation_date=trade_dt)
            print(f"择时信号: {timing_signal}")
    except Exception as e:
        print(f"❌ 快速测试失败: {e}")
    print('=' * 50)

if __name__ == "__main__":
    run_single_day()
